package com.da.result;

import com.da.exception.OllamaClientException;
import com.da.tools.Parse;
import com.da.tools.json.JsonMapperUtil;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.function.Function;

/**
 * 流式响应处理对象
 * <p>
 * 优化要点：
 * 1. 使用 try-with-resources 确保输出流安全关闭
 * 2. 明确指定 UTF-8 字符集避免编码问题
 * 3. 预定义解析策略减少条件判断
 * 4. 使用静态 JSON 解析器减少对象创建开销
 * 5. 增加缓冲区大小提升读取性能
 *
 * @author da
 * @time 2025/2/21 下午 5:42
 */
public class StreamResult implements Iterable<Result>, AutoCloseable {
    // 包装读取流（使用更大缓冲区提升读取性能）
    private static final int BUFFER_SIZE = 16384; // 16KB
    // HTTP 连接对象
    private final HttpURLConnection connection;

    // 缓冲读取器（使用明确字符集）
    private final BufferedReader reader;

    // 响应类型（决定数据解析策略）
    private final String type;

    /**
     * 流式响应构造函数
     *
     * @param url    请求地址
     * @param data   请求数据
     * @param method HTTP 方法（GET/POST）
     * @param type   响应类型（generate/chat）
     * @throws IOException 当发生 I/O 异常或 HTTP 状态非 200 时抛出
     */
    public StreamResult(String url, Map<String, Object> data, String method, String type) throws IOException {
        // 参数校验
        if (!"GET".equalsIgnoreCase(method) && !"POST".equalsIgnoreCase(method)) {
            throw new IllegalArgumentException("不支持的 HTTP 方法: " + method);
        }

        this.type = type;
        String json = Parse.toJson(data);

        // 创建 HTTP 连接
        connection = (HttpURLConnection) new URL(url).openConnection();
        connection.setRequestMethod(method.toUpperCase());

        // 处理 POST 请求
        if ("POST".equalsIgnoreCase(method)) {
            connection.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
            connection.setDoOutput(true);

            // 使用 try-with-resources 确保输出流安全关闭
            try (OutputStream os = connection.getOutputStream()) {
                os.write(json.getBytes(StandardCharsets.UTF_8));
                os.flush();
            }
        }

        // 校验 HTTP 响应状态
        int statusCode = connection.getResponseCode();
        if (statusCode != HttpURLConnection.HTTP_OK) {
            throw new IOException("HTTP 请求失败，状态码: " + statusCode);
        }

        // 使用带缓冲的读取器并指定字符集
        this.reader = new BufferedReader(
                new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8),
                BUFFER_SIZE
        );
    }

    /**
     * 关闭资源方法（实现 AutoCloseable 接口）
     * 关闭顺序：
     * 1. 先关闭输入流
     * 2. 再断开 HTTP 连接
     */
    @Override
    public void close() throws Exception {
        // 先关闭读取流
        try {
            reader.close();
        } finally {
            // 确保始终断开连接
            connection.disconnect();
        }
    }

    /**
     * 实现迭代器接口，用于流式处理响应数据
     * <p>
     * 优化要点：
     * 1. 预定义解析策略避免重复条件判断
     * 2. 优化异常处理流程
     * 3. 增加类型安全检查
     */
    @Override
    public Iterator<Result> iterator() {
        final Function<String, Result> parser = createParser();

        return new Iterator<Result>() {
            private String nextLine;
            private boolean closed;
            private boolean needsRead = true;  // 新增读取需求标记

            @Override
            public boolean hasNext() {
                if (closed) return false;

                try {
                    if (needsRead) {
                        nextLine = reader.readLine();
                        needsRead = false;
                    }
                    if (nextLine == null) {
                        closeResources();
                    }
                    return nextLine != null;
                } catch (IOException e) {
                    closeResources();
                    throw new OllamaClientException("流式读取失败", e);
                }
            }

            @Override
            public Result next() {
                if (!hasNext()) {
                    throw new NoSuchElementException("没有更多数据");
                }

                try {
                    return parser.apply(nextLine);
                } finally {
                    needsRead = true;  // 标记下次需要读取新行
                    nextLine = null;   // 清空当前行
                }
            }

            private void closeResources() {
                if (!closed) {
                    try {
                        closed = true;
                        StreamResult.this.close();
                    } catch (Exception e) {
                        throw new OllamaClientException("资源关闭失败", e);
                    }
                }
            }
        };
    }

    /**
     * 创建对应的数据解析策略
     *
     * @return 解析函数实例
     * @throws IllegalArgumentException 当遇到未知类型时抛出
     */
    private Function<String, Result> createParser() {
        switch (type) {
            case "generate":
                return line -> JsonMapperUtil.fromJson(line, MessageResult.class);
            case "chat":
                return line -> JsonMapperUtil.fromJson(line, ChatResult.class);
            default:
                throw new IllegalArgumentException("未知的响应类型: " + type);
        }
    }
}